from ..container.container import GuardContainer


class ComponentContainer(GuardContainer):
    components = {
        "train": {
            "dataloader": True,
            "transform": True,
            "model": True,
            "criteria": True,
            "optimizer": True,
            "data_evaluator": True,
            "loss_evaluator": True,
            "grad_evaluator": True,
            "dispatcher": False
        },
        "val": {
            "dataloader": True,
            "transform": True,
            "model": True,
            "criteria": True,
            "data_evaluator": True,
            "loss_evaluator": True,
            "dispatcher": False
        }
    }

    def __init__(self):
        super().__init__()
        self.put("runtime", {})
        self.put("extra", {})

    def verify(self, mode) -> bool:
        for component, required in self.components[mode].items():
            if required and not self.contains(component):
                return False

        return True

    @property
    def dataloader(self):
        return self.get("dataloader")

    @dataloader.setter
    def dataloader(self, dataloader):
        self.put("dataloader", dataloader)

    @property
    def transform(self):
        return self.get("transform")

    @transform.setter
    def transform(self, transform):
        self.put("transform", transform)

    @property
    def model(self):
        return self.get("model")

    @model.setter
    def model(self, model):
        self.put("model", model)

    @property
    def criteria(self):
        return self.get("criteria")

    @criteria.setter
    def criteria(self, criteria):
        self.put("criteria", criteria)

    @property
    def optimizer(self):
        return self.get("optimizer")

    @optimizer.setter
    def optimizer(self, optimizer):
        self.put("optimizer", optimizer)

    @property
    def data_evaluator(self):
        return self.get("data_evaluator")

    @data_evaluator.setter
    def data_evaluator(self, data_evaluator):
        self.put("data_evaluator", data_evaluator)

    @property
    def loss_evaluator(self):
        return self.get("loss_evaluator")

    @loss_evaluator.setter
    def loss_evaluator(self, loss_evaluator):
        self.put("loss_evaluator", loss_evaluator)

    @property
    def grad_evaluator(self):
        return self.get("grad_evaluator")

    @grad_evaluator.setter
    def grad_evaluator(self, grad_evaluator):
        self.put("grad_evaluator", grad_evaluator)

    @property
    def total_epoch(self):
        return self.get("runtime.total_epoch")

    @total_epoch.setter
    def total_epoch(self, total_epoch):
        self.put("runtime.total_epoch", total_epoch)

    @property
    def total_iter(self):
        return self.get("runtime.total_iter")

    @total_iter.setter
    def total_iter(self, total_iter):
        self.put("runtime.total_iter", total_iter)

    @property
    def epoch(self):
        return self.get("runtime.epoch")

    @epoch.setter
    def epoch(self, epoch):
        self.put("runtime.epoch", epoch)

    @property
    def iter(self):
        return self.get("runtime.iter")

    @iter.setter
    def iter(self, iter):
        self.put("runtime.iter", iter)

    @property
    def data(self):
        return self.get("runtime.data")

    @data.setter
    def data(self, data):
        self.put("runtime.data", data)

    @property
    def pred(self):
        return self.get("runtime.pred")

    @pred.setter
    def pred(self, pred):
        self.put("runtime.pred", pred)

    @property
    def loss(self):
        return self.get("runtime.loss")

    @loss.setter
    def loss(self, loss):
        self.put("runtime.loss", loss)

    @property
    def dispatcher(self):
        return self.get("dispatcher")

    @dispatcher.setter
    def dispatcher(self, dispatcher):
        self.put("dispatcher", dispatcher)

    @property
    def extra(self):
        return self.get("extra")

    def put_extra(self, key, val):
        self.put("extra." + key, val)

    def get_extra(self, key):
        return self.get("extra." + key)
